"""
@author: badmonkey
@software: PyCharm
@file: LL1.py
@time: 2021/4/14 上午10:34
"""

input_txt = "../Lexer/bubble-Lexer.txt"
output_doc_part = "bubble-"

# 文法定义，使用字典表示SNL语言的上下文无关文法
grammar = {
    # 起始符，表示程序的入口非终结符
    "S": "Program",
    # 产生式集合，定义了每个非终结符的推导规则
    "P": {
        "Program": {0: ["ProgramHead", "DeclarePart", "ProgramBody", "."]},  # 程序由头部、声明部分、主体和结束符"."组成
        "ProgramHead": {0: ["PROGRAM", "ProgramName"]},
        "ProgramName": {0: ["ID"]},
        "DeclarePart": {0: ["TypeDec", "VarDec", "ProcDec"]},
        "TypeDec": {0: ["ε"], 1: ["TypeDeclaration"]},
        "TypeDeclaration": {0: ["TYPE", "TypeDecList"]},
        "TypeDecList": {0: ["TypeId", "=", "TypeName", ";", "TypeDecMore"]},
        "TypeDecMore": {0: ["ε"], 1: ["TypeDecList"]},
        "TypeId": {0: ["ID"]},
        "TypeName": {0: ["BaseType"], 1: ["StructureType"], 2: ["ID"]},
        "BaseType": {0: ["INTEGER"], 1: ["CHAR"]},
        "StructureType": {0: ["ArrayType"], 1: ["RecType"]},
        "ArrayType": {0: ["ARRAY", "[", "Low", "..", "Top", "]", "OF", "BaseType"]},
        "Low": {0: ["INTC"]},
        "Top": {0: ["INTC"]},
        "RecType": {0: ["RECORD", "FieldDecList", "END"]},
        "FieldDecList": {
            0: ["BaseType", "IdList", ";", "FieldDecMore"],
            1: ["ArrayType", "IdList", ";", "FieldDecMore"],
        },
        "FieldDecMore": {0: ["ε"], 1: ["FieldDecList"]},
        "IdList": {0: ["ID", "IdMore"]},
        "IdMore": {0: ["ε"], 1: [",", "IdList"]},
        "VarDec": {0: ["ε"], 1: ["VarDeclaration"]},
        "VarDeclaration": {0: ["VAR", "VarDecList"]},
        "VarDecList": {0: ["TypeName", "VarIdList", ";", "VarDecMore"]},
        "VarDecMore": {0: ["ε"], 1: ["VarDecList"]},
        "VarIdList": {0: ["ID", "VarIdMore"]},
        "VarIdMore": {0: ["ε"], 1: [",", "VarIdList"]},
        "ProcDec": {0: ["ε"], 1: ["ProcDeclaration"]},
        "ProcDeclaration": {
            0: ["PROCEDURE", "ProcName", "(", "ParamList", ")", ";", "ProcDecPart", "ProcBody", "ProcDecMore"]
        },
        "ProcDecMore": {0: ["ε"], 1: ["ProcDeclaration"]},
        "ProcName": {0: ["ID"]},
        "ParamList": {0: ["ε"], 1: ["ParamDecList"]},
        "ParamDecList": {0: ["Param", "ParamMore"]},
        "ParamMore": {0: ["ε"], 1: [";", "ParamDecList"]},
        "Param": {0: ["TypeName", "FormList"], 1: ["VAR", "TypeName", "FormList"]},
        "FormList": {0: ["ID", "FidMore"]},
        "FidMore": {0: ["ε"], 1: [",", "FormList"]},
        "ProcDecPart": {0: ["DeclarePart"]},
        "ProcBody": {0: ["ProgramBody"]},
        "ProgramBody": {0: ["BEGIN", "StmList", "END"]},
        "StmList": {0: ["Stm", "StmMore"]},
        "StmMore": {0: ["ε"], 1: [";", "StmList"]},
        "Stm": {  # 语句可以是条件语句、循环语句、输入语句、输出语句、返回语句或赋值/调用语句
            0: ["ConditionalStm"],
            1: ["LoopStm"],
            2: ["InputStm"],
            3: ["OutputStm"],
            4: ["ReturnStm"],
            5: ["ID", "AssCall"],
        },
        "AssCall": {0: ["AssignmentRest"], 1: ["CallStmRest"]},
        "AssignmentRest": {0: ["VariMore", ":=", "Exp"]},
        "ConditionalStm": {
            0: ["IF", "RelExp", "THEN", "StmList", "ELSE", "StmList", "FI"]
        },
        "LoopStm": {0: ["WHILE", "RelExp", "DO", "StmList", "ENDWH"]},
        "InputStm": {0: ["READ", "(", "Invar", ")"]},
        "Invar": {0: ["ID"]},
        "OutputStm": {0: ["WRITE", "(", "Exp", ")"]},
        "ReturnStm": {0: ["RETURN", "(", "Exp", ")"]},
        "CallStmRest": {0: ["(", "ActParamList", ")"]},
        "ActParamList": {0: ["ε"], 1: ["Exp", "ActParamMore"]},
        "ActParamMore": {0: ["ε"], 1: [",", "ActParamList"]},
        "RelExp": {0: ["Exp", "OtherRelE"]},
        "OtherRelE": {0: ["CmpOp", "Exp"]},
        "Exp": {0: ["Term", "OtherTerm"]},
        "OtherTerm": {0: ["ε"], 1: ["AddOp", "Exp"]},
        "Term": {0: ["Factor", "OtherFactor"]},
        "OtherFactor": {0: ["ε"], 1: ["MultOp", "Term"]},
        "Factor": {0: ["(", "Exp", ")"], 1: ["INTC"], 2: ["Variable"]},
        "Variable": {0: ["ID", "VariMore"]},
        "VariMore": {0: ["ε"], 1: ["[", "Exp", "]"], 2: [".", "FieldVar"]},
        "FieldVar": {0: ["ID", "FieldVarMore"]},
        "FieldVarMore": {0: ["ε"], 1: ["[", "Exp", "]"]},
        "CmpOp": {0: ["<"], 1: ["="]},
        "AddOp": {0: ["+"], 1: ["-"]},
        "MultOp": {0: ["*"], 1: ["/"]},
    },
    # 非终结符集合，文法中所有非终结符的列表
    "VN": [
        "ActParamList", "ActParamMore", "AddOp", "ArrayType", "AssCall", "AssignmentRest",
        "BaseType", "CallStmRest", "CmpOp", "ConditionalStm", "DeclarePart", "Exp",
        "Factor", "FidMore", "FieldDecList", "FieldDecMore", "FieldVar", "FieldVarMore",
        "FormList", "IdList", "IdMore", "InputStm", "Invar", "LoopStm", "Low", "MultOp",
        "OtherFactor", "OtherRelE", "OtherTerm", "OutputStm", "Param", "ParamDecList",
        "ParamList", "ParamMore", "ProcBody", "ProcDec", "ProcDecMore", "ProcDecPart",
        "ProcDeclaration", "ProcName", "Program", "ProgramBody", "ProgramHead", "ProgramName",
        "RecType", "RelExp", "ReturnStm", "Stm", "StmList", "StmMore", "StructureType",
        "Term", "Top", "TypeDec", "TypeDecList", "TypeDecMore", "TypeDeclaration", "TypeId",
        "TypeName", "VarDec", "VarDecList", "VarDecMore", "VarDeclaration", "VarIdList",
        "VarIdMore", "VariMore", "Variable",
    ],
    # 终结符集合，文法中所有终结符的列表（关键字、运算符、标点等）
    "VT": [
        "(", ")", "*", "+", ",", "-", ".", "..", "/", ":=", ";", "<", "=", "ARRAY", "BEGIN",
        "CHAR", "DO", "ELSE", "END", "ENDWH", "FI", "ID", "IF", "INTC", "INTEGER", "OF",
        "PROCEDURE", "PROGRAM", "READ", "RECORD", "RETURN", "THEN", "TYPE", "VAR", "WHILE",
        "WRITE", "[", "]",
    ],
    # 空符号，表示空产生式
    "NN": ["ε"],
}


# 计算FIRST集合的函数
def firstSet(grammar, leftSymbol):
    """
    计算给定符号的FIRST集合，即从该符号开始推导出的所有可能首终结符集合
    :param grammar: 文法规则
    :param leftSymbol: 需要计算FIRST集合的符号（可以是终结符或非终结符）
    :return: FIRST集合（set类型）
    """
    result = set()
    # 如果是终结符或空符号，直接加入结果集合
    if leftSymbol in grammar["VT"] or leftSymbol == "ε":
        result.add(leftSymbol)
        return result
    # 获取该非终结符的所有产生式
    sons = grammar["P"][leftSymbol]
    # 遍历每个产生式
    for i in range(len(sons)):
        son = sons[i]  # 当前产生式右部
        # 遍历产生式右部的每个符号
        for grandSon in son:
            grandSonFirst = firstSet(grammar, grandSon)  # 递归计算子符号的FIRST集合
            result = result.union(grandSonFirst - {"ε"})  # 加入非空符号
            if "ε" not in grandSonFirst:  # 如果子符号不能推导出空，则停止
                break
        # 如果产生式右部的所有符号都能推导出空，则将空加入结果
        if all("ε" in firstSet(grammar, s) for s in son):
            result.add("ε")
    return result


# 计算FOLLOW集合的函数
def followSet(grammar):
    """
    计算所有非终结符的FOLLOW集合，即可能紧跟在该非终结符后的终结符集合
    :param grammar: 文法规则
    :return: 所有非终结符的FOLLOW集合（字典类型）
    """
    # 初始化所有非终结符的FOLLOW集合为空
    allFollowSet = {}
    for vn in grammar["VN"]:
        allFollowSet[vn] = set()
    # 起始符的FOLLOW集合初始化包含结束符"#"
    allFollowSet[grammar["S"]].add("#")

    done = False
    # 使用迭代法计算FOLLOW集合，直到不再变化
    while not done:
        preStatus = [len(allFollowSet[vn]) for vn in grammar["VN"]]  # 记录当前状态
        # 遍历所有产生式
        for p in grammar["P"]:
            sons = grammar["P"][p]
            length = len(sons)
            for i in range(length):
                son = sons[i]  # 当前产生式右部
                sonNum = len(son)
                # 遍历产生式右部的每个符号
                for j in range(sonNum):
                    grandSon = son[j]
                    if grandSon not in grammar["VN"]:  # 跳过终结符
                        continue
                    else:
                        # 如果是产生式最后一个符号，将左部的FOLLOW集合加入
                        if j == sonNum - 1:
                            allFollowSet[grandSon] = allFollowSet[grandSon].union(allFollowSet[p])
                        else:
                            k = j + 1
                            backSon = son[k]
                            backSonFirst = firstSet(grammar, backSon)  # 计算后继符号的FIRST集合
                            # 如果后继符号能推导出空，则继续向后看
                            while k < sonNum and "ε" in backSonFirst:
                                backSonFirst.remove("ε")
                                allFollowSet[grandSon] = allFollowSet[grandSon].union(backSonFirst)
                                k += 1
                                if k == sonNum:
                                    allFollowSet[grandSon] = allFollowSet[grandSon].union(allFollowSet[p])
                                    break
                                backSon = son[k]
                                backSonFirst = firstSet(grammar, backSon)
                            allFollowSet[grandSon] = allFollowSet[grandSon].union(backSonFirst)

        curStatus = [len(allFollowSet[vn]) for vn in grammar["VN"]]  # 记录更新后的状态
        if curStatus == preStatus:  # 如果状态未变化，则结束
            done = True
    return allFollowSet


# 构建LL(1)分析表的函数
def ll1Table(grammar):
    """
    构建LL(1)分析表，用于预测解析
    """
    first = {}  # 存储所有非终结符的FIRST集合
    table = {}  # LL(1)分析表
    # 初始化分析表，所有项默认值为"error"
    for vn in grammar["VN"]:
        first[vn] = firstSet(grammar, vn)  # 求vn这个非终极符的first集
        table[vn] = {}
        for vt in grammar["VT"]:
            table[vn][vt] = "error"
    follow = followSet(grammar)  # 计算FOLLOW集合

    # 填充LL1表
    for p in grammar["P"]:
        # p是产生式的左部(非终结符)
        sons = grammar["P"][p]
        # p的产生式右部有多个，用sons[i]表示其中每一个
        for i in range(len(sons)):
            son = sons[i]  # 产生式右部
            head = son[0]  # 产生式右部的第一个符号
            if head in grammar["VT"]:  # 如果是终结符，直接将产生式填入表中(等效于predict(p)=head)
                table[p][head] = son
            else:
                possible = set()  # predict(p)的候选集合
                flag = True
                for element in son:
                    # 遍历产生式右部的每个符号（注意下面可以break出去）
                    if element in grammar["VN"]:  # 非终结符
                        tmp = first[element].copy()
                        if "ε" in tmp:  # 如果能推导出空，则将first集加入FOLLOW集合
                            tmp.remove("ε")
                            possible = possible.union(tmp)
                        else:
                            # predict(p) = first(element)
                            possible = possible.union(tmp)
                            flag = False
                            break
                    elif element == "ε":  # 空产生式，加入FOLLOW集合
                        # predict(p) = follow(p)
                        possible = possible.union(follow[p])
                    elif element in grammar["VT"]:  # 终结符
                        possible.add(element)
                        flag = False
                        break
                if flag:
                    # 如果产生式右部的所有符号都能推导出空，则加入FOLLOW集合
                    # predict(p) = first(son) ∪ follow(p)
                    possible = possible.union(follow[p])
                for poss in possible:
                    # predict(p)=poss
                    table[p][poss] = son  # 填充分析表

    # 调试输出：打印VarDec的表项
    print("VarDec 的表项：")
    for vt in grammar["VT"]:
        if table["VarDec"][vt] != "error":
            print(f"  {vt}: {table['VarDec'][vt]}")

    # 将FIRST集合写入文件
    with open(f"{output_doc_part}first.txt", "w") as f:
        for vn in first:
            f.write(f"{vn}\t{first[vn]}\n")
    # 将FOLLOW集合写入文件
    with open(f"{output_doc_part}follow.txt", "w") as f:
        for vn in follow:
            f.write(f"{vn}\t{follow[vn]}\n")

    return table


from Parser.AST import *  # 导入AST相关类（假设定义了AstNode和AstNodeEncoder）


# 生成抽象语法树（AST）的函数
def generateAST(tokens):
    """
    根据词法分析的标记生成抽象语法树
    :param tokens: 词法分析器输出的标记列表
    :return: AST的根节点
    """
    tokenStack = []  # 用于跟踪输入标记的栈（输入的原始token流）
    table = ll1Table(grammar)  # 构建LL(1)分析表

    # 把每一个token的类型填入tokenStack
    for token in tokens:
        token = eval(token.strip())  # 将字符串形式的标记转换为Python对象
        tokenType = token[0]
        tokenStack.append(tokenType)

    stack = [grammar["S"]]  # 解析栈，初始为起始符（自顶向下变化的栈）
    root = AstNode(grammar["S"])  # 创建AST根节点
    current = root  # 当前操作的AST节点

    i = 0  # 当前处理的token索引
    while i < len(tokens) and stack:  # 当还有token未处理且栈不为空时继续
        token = eval(tokens[i].strip())
        tokenType = token[0]  # 标记类型
        tokenVal = token[1]  # 标记值

        top = stack[-1]  # 查看栈顶符号
        print(f"Step {i}: Stack={stack}, Top={top}, Token={tokenType}")  # 调试信息

        if top == "ε":  # 处理空产生式
            stack.pop()
            current.insertChild(AstNode("ε", "ε"))  # 在AST中添加空节点
            current = current.step()  # 移动到下未完全扩展完成的兄弟节点
            continue

        if top == tokenType:  # 栈顶终结符与当前token匹配
            stack.pop()  # 弹出栈顶
            tokenStack.pop(0)  # 移除已处理的token
            current.tokenType = tokenType  # 设置AST节点的类型和值
            current.tokenVal = tokenVal
            current = current.step()  # 移动到下一个兄弟节点
            i += 1
        elif top in grammar["VT"]:  # 栈顶是终结符但不匹配
            print(f"错误：期望终结符 '{top}' 但得到 '{tokenType}' 在 token {token}")
            i += 1  # 跳过当前token尝试恢复
        else:  # 栈顶是非终结符
            try:
                choice = table[top][tokenType]  # 查询分析表
                if choice == "error":  # 没有匹配的产生式
                    print(f"错误：'{top}' 对于输入 '{tokenType}' 无有效产生式，在 token {token}")
                    i += 1  # 跳过token尝试恢复
                    continue
                stack.pop()  # 弹出非终结符
                # 将产生式右部逆序压入栈
                for symbol in choice[::-1]:
                    stack.append(symbol)
                # 为产生式创建AST子节点
                for symbol in choice:
                    current.insertChild(AstNode(symbol))
                current = current.child[0]  # 移动到第一个子节点
            except KeyError:  # 分析表中无对应项
                print(f"错误：'{top}' 对于输入 '{tokenType}' 无表项，在 token {token}")
                i += 1  # 跳过token尝试恢复

    # 检查解析是否完成
    if stack:
        print(f"错误：解析未完成，栈中剩余：{stack}")
    if i < len(tokens):
        print(f"错误：未处理的 token：{tokens[i:]}")

    # 将AST以文本格式输出到文件
    with open(f"{output_doc_part}ast.txt", "w", encoding="utf-8") as ast_file:
        root.dump(file=ast_file)

    # 将AST以JSON格式输出到文件
    with open(f"{output_doc_part}ast_json.txt", "w", encoding="utf-8") as json_ast_file:
        json.dump(root, json_ast_file, cls=AstNodeEncoder)

    return root


# 可视化AST的函数
def display(root):
    """
    使用Graphviz生成并显示抽象语法树的图形化表示
    :param root: AST的根节点
    """
    current = root
    from graphviz import Digraph

    graph = Digraph(name=f"{output_doc_part}ast", format="png")  # 创建Graphviz对象
    stack = []
    stack.append(current)  # 使用栈进行树的遍历
    while stack:
        node = stack.pop(0)  # 取出栈顶节点
        color = "black"  # 默认节点颜色
        name = "node" + str(node.getId())  # 节点唯一标识
        label = node.getTokenType()  # 节点标签
        if label == "ε":  # 空节点用黄色表示
            color = "yellow"
        elif label in grammar["VT"]:  # 终结符用红色表示
            color = "red"
            if label in ["INTC", "ID"]:  # 对于整数常量和标识符，显示其值
                label = label + "\n" + str(node.getTokenVal())

        graph.node(name=name, label=label, color=color)  # 添加节点
        for child in node.child:  # 添加子节点并创建边
            sonName = "node" + str(child.getId())
            stack.append(child)
            graph.edge(name, sonName)
    graph.render(f"{output_doc_part}ast", view=True)  # 渲染并显示AST图形


# 主程序入口
if __name__ == "__main__":
    # 读取词法分析的输出文件
    tokens = open(input_txt, "r").readlines()
    # 生成AST
    root = generateAST(tokens)
    # 可视化AST
    display(root)
